import sys
import shutil
import os
import cv2
import yaml
import torch
from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam import GradCAMPlusPlus, GradCAM, XGradCAM
from utils.general import xywh2xyxy
from utils.datasets import letterbox
from utils.torch_utils import intersect_dicts
from models.yolo import Model
from PIL import Image
from tqdm import trange
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import warnings
warnings.filterwarnings('ignore')
warnings.simplefilter('ignore')
np.random.seed(0)

sys.path.append(os.getcwd())


class yolov7_heatmap:
    def __init__(self, weight, cfg, device, method, layer, backward_type, conf_threshold, ratio):
        device = torch.device(device)
        ckpt = torch.load(weight)
        model_names = ckpt['model'].names
        # checkpoint state_dict as FP32
        csd = ckpt['model'].float().state_dict()
        model = Model(cfg, ch=3, nc=len(model_names)).to(device)
        csd = intersect_dicts(csd, model.state_dict(),
                              exclude=['anchor'])  # intersect
        model.load_state_dict(csd, strict=False)  # load
        model.eval()
        print(f'Transferred {len(csd)}/{len(model.state_dict())} items')

        target_layers = [eval(layer)]
        method = eval(method)

        colors = np.random.uniform(0, 255, size=(
            len(model_names), 3)).astype(np.int)
        self.__dict__.update(locals())

    def post_process(self, result):
        boxes_ = result[0][..., :4]
        logits_ = []
        for data in result[1]:
            bs, n, w, h, _ = data.size()
            logits_.append(data.reshape((bs, n * w * h, _)))
        logits_ = torch.cat(logits_, dim=1)[..., 4:]
        sorted, indices = torch.sort(logits_[..., 0], descending=True)
        logits_ = logits_[0][indices[0]]
        logits_[:, 0] = torch.sigmoid(logits_[:, 0])
        return logits_, xywh2xyxy(boxes_[0][indices[0]]).cpu().detach().numpy()

    def draw_detections(self, box, color, name, img):
        xmin, ymin, xmax, ymax = list(map(int, list(box)))
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax),
                      tuple(int(x) for x in color), 2)
        cv2.putText(img, str(name), (xmin, ymin - 5), cv2.FONT_HERSHEY_SIMPLEX,
                    0.8, tuple(int(x) for x in color), 2, lineType=cv2.LINE_AA)
        return img

    def __call__(self, img_path, save_path):
        # remove dir if exist
        if os.path.exists(save_path):
            shutil.rmtree(save_path)
        # make dir if not exist
        os.makedirs(save_path, exist_ok=True)

        # img process
        img = cv2.imread(img_path)
        # # ----------------------------读入差分----------------------------
        # ref_image = cv2.imread("inference/background/0.jpg")
        # diff_image = cv2.absdiff(ref_image, img)
        # diff_image = cv2.cvtColor(
        #     diff_image, cv2.COLOR_BGR2GRAY)[..., None]
        # img = np.concatenate((img, diff_image), axis=-1)

        # # ----------------------------结束差分----------------------------

        img = letterbox(img)[0]
        img = img[:, :, ::-1]
        img = np.float32(img) / 255.0

        tensor = torch.from_numpy(np.transpose(
            img, axes=[2, 0, 1])).unsqueeze(0).to(self.device)

        # init ActivationsAndGradients
        grads = ActivationsAndGradients(
            self.model, self.target_layers, reshape_transform=None)

        # get ActivationsAndResult
        result = grads(tensor)
        activations = grads.activations[0].cpu().detach().numpy()

        # postprocess to yolo output
        post_result, post_boxes = self.post_process(result)
        for i in trange(int(post_result.size(0) * self.ratio)):
            if post_result[i][0] < self.conf_threshold:
                break

            self.model.zero_grad()
            if self.backward_type == 'conf':
                post_result[i, 0].backward(retain_graph=True)
            else:
                # get max probability for this prediction
                score = post_result[i, 1:].max()
                score.backward(retain_graph=True)

            # process heatmap
            gradients = grads.gradients[0]
            b, k, u, v = gradients.size()
            weights = self.method.get_cam_weights(
                self.method, None, None, None, activations, gradients.detach().numpy())
            weights = weights.reshape((b, k, 1, 1))
            saliency_map = np.sum(weights * activations, axis=1)
            saliency_map = np.squeeze(np.maximum(saliency_map, 0))
            saliency_map = cv2.resize(
                saliency_map, (tensor.size(3), tensor.size(2)))
            saliency_map_min, saliency_map_max = saliency_map.min(), saliency_map.max()
            if (saliency_map_max - saliency_map_min) == 0:
                continue
            saliency_map = (saliency_map - saliency_map_min) / \
                (saliency_map_max - saliency_map_min)

            # add heatmap and box to image
            cam_image = show_cam_on_image(
                img.copy(), saliency_map, use_rgb=True)
            cam_image = self.draw_detections(post_boxes[i], self.colors[int(post_result[i, 1:].argmax(
            ))], f'{self.model_names[int(post_result[i, 1:].argmax())]} {post_result[i][0]:.2f}', cam_image)
            cam_image = Image.fromarray(cam_image)
            cam_image.save(f'{save_path}/{i}.png')


def get_params():
    params = {
        'weight': 'modelweight/final.pt',
        'cfg': 'cfg/training/myyolov7.yaml',
        'device': 'cuda:0',
        'method': 'GradCAM',  # GradCAMPlusPlus, GradCAM, XGradCAM
        'layer': 'model.model[-2]',
        'backward_type': 'class',  # class or conf
        'conf_threshold': 0.6,  # 0.6
        'ratio': 0.02  # 0.02-0.1
    }
    return params


if __name__ == '__main__':
    model = yolov7_heatmap(**get_params())
    model('inference/elevator/999_97.jpg', 'result')
